from random import randint
import time
IV = [0x7380166F, 0x4914B2B9, 0x172442D7, 0xDA8A0600, 0xA96F30BC, 0x163138AA, 0xE38DEE4D, 0xB0FB0E4E]
T = [0x79cc4519, 0x7a879d8a]

class SM3_unit:
   def Rol(X, i):
       i = i % 32
       return ((X << i) & 0xFFFFFFFF) | ((X & 0xFFFFFFFF) >> (32 - i))

   def F(a, b, c, j):
       if j >= 0 and j <= 15:
           return a ^ b ^ c
       else:
           return ((a & b) | (a & c) | (b & c))

   def G(a, b, c, j):
       if j >= 0 and j <= 15:
           return a ^ b ^ c
       else:
           return ((a & b) | (~a & c))

   def P0(a):
       return a ^ SM3_unit.Rol(a, 9) ^ SM3_unit.Rol(a, 17)

   def P1(a):
       return a ^ SM3_unit.Rol(a, 15) ^ SM3_unit.Rol(a, 23)

   def T_(j):
       if j >= 0 and j <= 15:
           return T[0]
       else:
           return T[1]


def Fill(info):
    m = bin(int(info, 16))[2:]
    if len(m) != len(info) * 4:
        m = '0' * (len(info) * 4 - len(m)) + m
    l = len(m)
    binary = '0' * (64 - len(bin(l)[2:])) + bin(l)[2:]
    m = m + '1'
    if len(m) % 512 > 448:
        temp=(512 - len(m) % 512 + 448)
        m = m + '0' * temp + binary
    else:
        temp_=(448 - len(m) % 512)
        m = m + '0' * temp_ + binary
    m = hex(int(m, 2))[2:]
    return m


def Conbind(m):
    n = len(m) / 128
    M = []
    for i in range(int(n)):
        M.append(m[0 + 128 * i:128 + 128 * i])
    return M


def Extend(M, n):
    W = []
    W_ = []
    for j in range(16):
        W.append(int(M[n][0 + 8 * j:8 + 8 * j], 16))
    for j in range(16, 68):
        W.append(SM3_unit.P1(W[j - 16] ^ W[j - 9] ^ SM3_unit.Rol(W[j - 3], 15)) ^ SM3_unit.Rol(W[j - 13], 7) ^ W[j - 6])
    for j in range(64):
        W_.append(W[j] ^ W[j + 4])
    Wstr = ''
    W_str = ''
    for x in W:
        Wstr += (hex(x)[2:] + ' ')
    for x in W_:
        W_str += (hex(x)[2:] + ' ')
    return W, W_


def CF(V, M, i):
    A, B, C, D, E, F, G, H = V[i]
    W, W_ = Extend(M, i)
    for j in range(64):
        SS1 = SM3_unit.Rol((SM3_unit.Rol(A, 12) + E + SM3_unit.Rol(SM3_unit.T_(j), j % 32)) % (2 ** 32), 7)
        SS2 = SS1 ^ SM3_unit.Rol(A, 12)
        TT1 = (SM3_unit.F(A, B, C, j) + D + SS2 + W_[j]) % (2 ** 32)
        TT2 = (SM3_unit.G(E, F, G, j) + H + SS1 + W[j]) % (2 ** 32)
        D = C
        C = SM3_unit.Rol(B, 9)
        B = A
        A = TT1
        H = G
        G = SM3_unit.Rol(F, 19)
        F = E
        E = SM3_unit.P0(TT2)
    a, b, c, d, e, f, g, h = V[i]
    V_ = [a ^ A, b ^ B, c ^ C, d ^ D, e ^ E, f ^ F, g ^ G, h ^ H]
    return V_


def Iterate(M):
    n = len(M)
    V = []
    V.append(IV)
    for i in range(n):
        V.append(CF(V, M, i))
    return V[n]


def SM3(message):
    m = Fill(message)  
    M = Conbind(m)  
    Vn = Iterate(M)  
    res = ''
    for x in Vn:
            res += hex(x)[2:]
    return res


def Rho(info,num):               
    x = hex(randint(0, 2**(info+1)-1))[2:]#从info bit字节随机生成一个十六进制数
    x_a = SM3(x)               
    x_b = SM3(x_a)              
    i = 1
    while x_a[:num] != x_b[:num]:
        i += 1
        x_a = SM3(x_a)              
        x_b = SM3(SM3(x_b))     
    x_b = x_a
    x_a = x            
    for j in range(i):
        if SM3(x_a)[:num] == SM3(x_b)[:num]:#找到一个m满足SM3(m)==SM3(SM3(m)),返回找到的num比特碰撞，以及两个碰撞的原信息。
            return SM3(x_a)[:num], x_a, x_b
        else:
            x_a = SM3(x_a)
            x_b = SM3(x_b)

while(True):
   info=int(input("Enter the message you want to counterfeit:"))
   bit=int(input("Please enter the collision bit length:"))
   t1=time.time()
   col, m1, m2 = Rho(info,bit)
   t2=time.time()
   print("Fing crash! \nThe first message is {}.\nThe second message is {}.\nThe hexadecimal value is {}.".format(m1,m2,col))
   print("Cost time:",t2-t1)
